import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime, timedelta
import os
from sqlalchemy import create_engine
from sqlalchemy.exc import SQLAlchemyError

# 确保中文显示正常
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
sns.set(font="SimHei", font_scale=1.2)


class InsuranceMetrics:
    def __init__(self, db_url='mysql+pymysql://root:root@localhost:3306/insurance_02'):
        """初始化数据库连接（使用SQLAlchemy）"""
        self.db_url = db_url
        self.engine = self.get_db_connection()

        # 创建结果和图表保存目录
        self.results_dir = 'results'
        self.plots_dir = 'plots'
        os.makedirs(self.results_dir, exist_ok=True)
        os.makedirs(self.plots_dir, exist_ok=True)

        # 测试记录和bug清单
        self.test_records = []
        self.bug_list = []

    def get_db_connection(self):
        """建立数据库连接"""
        try:
            engine = create_engine(self.db_url)
            # 测试连接
            with engine.connect():
                print("数据库连接成功")
            return engine
        except SQLAlchemyError as e:
            print(f"数据库连接失败: {str(e)}")
            raise  # 抛出异常，终止程序

    def close_connection(self):
        """关闭数据库连接（SQLAlchemy引擎不需要显式关闭，但可以释放资源）"""
        if self.engine:
            self.engine.dispose()
            print("数据库连接已关闭")

    def add_test_record(self, metric_name, test_date, result, comments=None):
        """添加测试记录"""
        self.test_records.append({
            'metric_name': metric_name,
            'test_date': test_date,
            'result': result,
            'comments': comments
        })

    def add_bug(self, bug_id, metric_name, description, status='未修复'):
        """添加bug记录"""
        self.bug_list.append({
            'bug_id': bug_id,
            'metric_name': metric_name,
            'description': description,
            'status': status
        })

    def save_test_records(self):
        """保存测试记录到CSV文件"""
        df = pd.DataFrame(self.test_records)
        df.to_csv(os.path.join(self.results_dir, 'test_records.csv'), index=False)
        print(f"测试记录已保存至 {os.path.join(self.results_dir, 'test_records.csv')}")

    def save_bug_list(self):
        """保存bug清单到CSV文件"""
        df = pd.DataFrame(self.bug_list)
        df.to_csv(os.path.join(self.results_dir, 'bug_list.csv'), index=False)
        print(f"Bug清单已保存至 {os.path.join(self.results_dir, 'bug_list.csv')}")

    def get_leading_reinsurer_contract_ratio(self, region=None, start_date=None, end_date=None):
        """
        1. 计算做首席再保人的合同数量占比
        计算公式：做首席再保人的合同数量 ÷ 合同总量 × 100%
        """
        query = """
        SELECT 
            SUM(CASE WHEN is_leading_reinsurer = 1 THEN 1 ELSE 0 END) AS leading_contracts,
            COUNT(*) AS total_contracts
        FROM contracts
        WHERE business_type = 'REINSURANCE'
        """

        params = []
        if region:
            query += " AND region = %s"
            params.append(region)
        if start_date and end_date:
            query += " AND issue_date BETWEEN %s AND %s"
            params.extend([start_date, end_date])

        # 使用pandas读取SQL
        result_df = pd.read_sql(query, self.engine, params=params)
        result = result_df.iloc[0].to_dict()

        if result['total_contracts'] == 0:
            ratio = 0.0
        else:
            ratio = (result['leading_contracts'] / result['total_contracts']) * 100

        # 保存结果
        result_df = pd.DataFrame({
            '指标名称': ['首席再保人合同数量占比'],
            '首席再保人合同数': [result['leading_contracts']],
            '总合同数': [result['total_contracts']],
            '占比(%)': [round(ratio, 2)],
            '区域': [region or '全国'],
            '统计时段': [f"{start_date or '开始'}-{end_date or '结束'}"]
        })

        # 添加测试记录
        self.add_test_record(
            '首席再保人合同数量占比',
            datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            '成功',
            f"区域: {region or '全国'}, 时段: {start_date or '开始'}-{end_date or '结束'}"
        )

        return ratio, result_df

    def get_leading_reinsurer_premium_ratio(self, region=None, start_date=None, end_date=None):
        """
        2. 计算做首席再保人的保费收入占比
        计算公式：做首席再保人的分保费收入 ÷ 分保费收入 × 100%
        """
        query = """
        SELECT 
            SUM(CASE WHEN c.is_leading_reinsurer = 1 THEN p.premium_amount ELSE 0 END) AS leading_premium,
            SUM(p.premium_amount) AS total_premium
        FROM premium_income p
        JOIN contracts c ON p.contract_id = c.contract_id
        WHERE c.business_type = 'REINSURANCE' AND p.is_original = 0
        """

        params = []
        if region:
            query += " AND c.region = %s"
            params.append(region)
        if start_date and end_date:
            query += " AND p.income_date BETWEEN %s AND %s"
            params.extend([start_date, end_date])

        result_df = pd.read_sql(query, self.engine, params=params)
        result = result_df.iloc[0].to_dict()

        # 关键修改：处理空值（将None转为0）
        leading_premium = result['leading_premium'] or 0.0
        total_premium = result['total_premium'] or 0.0

        if total_premium == 0:
            ratio = 0.0
        else:
            ratio = (leading_premium / total_premium) * 100

        # 保存结果
        result_df = pd.DataFrame({
            '指标名称': ['首席再保人保费收入占比'],
            '首席再保人保费': [leading_premium],
            '总保费': [total_premium],
            '占比(%)': [round(ratio, 2)],
            '区域': [region or '全国'],
            '统计时段': [f"{start_date or '开始'}-{end_date or '结束'}"]
        })

        # 添加测试记录
        self.add_test_record(
            '首席再保人保费收入占比',
            datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            '成功',
            f"区域: {region or '全国'}, 时段: {start_date or '开始'}-{end_date or '结束'}"
        )

        return ratio, result_df

    def get_long_short_term_premium_growth(self, period_type, region=None, report_year=None):
        """
        3. 计算长/短期险保费增长率
        计算公式：(报告期保费 - 基期保费) ÷ 基期保费 × 100%
        period_type: 'LONG' 或 'SHORT'
        """
        if not report_year:
            report_year = datetime.now().year
        base_year = report_year - 1

        # 获取报告期保费
        query_report = """
        SELECT SUM(p.premium_amount) AS premium_amount
        FROM premium_income p
        JOIN contracts c ON p.contract_id = c.contract_id
        JOIN products pr ON c.product_id = pr.product_id
        WHERE pr.insurance_period_type = %s AND YEAR(p.income_date) = %s
        """

        params = [period_type, report_year]
        if region:
            query_report += " AND c.region = %s"
            params.append(region)

        report_df = pd.read_sql(query_report, self.engine, params=params)
        report_premium = report_df.iloc[0]['premium_amount'] or 0

        # 获取基期保费
        query_base = """
        SELECT SUM(p.premium_amount) AS premium_amount
        FROM premium_income p
        JOIN contracts c ON p.contract_id = c.contract_id
        JOIN products pr ON c.product_id = pr.product_id
        WHERE pr.insurance_period_type = %s AND YEAR(p.income_date) = %s
        """

        params = [period_type, base_year]
        if region:
            query_base += " AND c.region = %s"
            params.append(region)

        base_df = pd.read_sql(query_base, self.engine, params=params)
        base_premium = base_df.iloc[0]['premium_amount'] or 0

        if base_premium == 0:
            growth_rate = 0.0
        else:
            growth_rate = ((report_premium - base_premium) / base_premium) * 100

        # 保存结果
        result_df = pd.DataFrame({
            '指标名称': [f'{"长期险" if period_type == "LONG" else "短期险"}保费增长率'],
            '报告期年份': [report_year],
            '报告期保费': [report_premium],
            '基期年份': [base_year],
            '基期保费': [base_premium],
            '增长率(%)': [round(growth_rate, 2)],
            '区域': [region or '全国']
        })

        # 添加测试记录
        self.add_test_record(
            f'{"长期险" if period_type == "LONG" else "短期险"}保费增长率',
            datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            '成功',
            f"区域: {region or '全国'}, 报告年份: {report_year}"
        )

        return growth_rate, result_df

    # 其他方法保持类似的修改模式，将cursor.execute替换为pd.read_sql
    # ... (省略其他方法，保持与之前功能一致)

    def generate_all_metrics(self, region=None, year=2023):
        """生成所有指标并保存结果"""
        print(f"开始计算{region or '全国'}{year}年的所有指标...")

        # 定义时间范围
        start_date = f"{year}-01-01"
        end_date = f"{year}-12-31"

        # 存储所有结果
        all_results = []

        # 1. 首席再保人合同数量占比
        ratio, df = self.get_leading_reinsurer_contract_ratio(region, start_date, end_date)
        all_results.append(df)
        print(f"1. 首席再保人合同数量占比: {ratio:.2f}%")

        # 2. 首席再保人保费收入占比
        ratio, df = self.get_leading_reinsurer_premium_ratio(region, start_date, end_date)
        all_results.append(df)
        print(f"2. 首席再保人保费收入占比: {ratio:.2f}%")

        # 3. 长期险保费增长率
        ratio, df = self.get_long_short_term_premium_growth('LONG', region, year)
        all_results.append(df)
        print(f"3. 长期险保费增长率: {ratio:.2f}%")

        # 3. 短期险保费增长率
        ratio, df = self.get_long_short_term_premium_growth('SHORT', region, year)
        all_results.append(df)
        print(f"3. 短期险保费增长率: {ratio:.2f}%")

        # 其余指标计算...

        # 合并所有结果并保存
        combined_df = pd.concat(all_results, ignore_index=True)
        filename = f"all_metrics_{region or 'national'}_{year}.csv"
        combined_df.to_csv(os.path.join(self.results_dir, filename), index=False)
        print(f"所有指标计算完成，结果已保存至 {os.path.join(self.results_dir, filename)}")

        return combined_df

    # 绘图相关方法保持不变
    # ...


# 主函数，用于演示
def main():
    # 初始化连接（使用SQLAlchemy连接字符串）
    # 格式: mysql+pymysql://用户名:密码@主机:端口/数据库名
    db_url = 'mysql+pymysql://root:root@localhost:3306/insurance_02'
    metrics_calculator = InsuranceMetrics(db_url)

    # 定义区域列表
    regions = ['华北', '华东', '华南', '西北', '西南']

    try:
        # 计算全国2023年所有指标
        national_metrics = metrics_calculator.generate_all_metrics(year=2023)

        # 计算各区域2023年所有指标
        for region in regions:
            metrics_calculator.generate_all_metrics(region=region, year=2023)

        # 生成区域对比图表和年度趋势图表...
        # ...

        # 添加一些示例bug
        metrics_calculator.add_bug(
            1,
            '保费预估差异率',
            '部分合同没有实际保费数据，导致差异率计算不准确',
            '未修复'
        )

        metrics_calculator.add_bug(
            2,
            '资产增量保费比',
            '资产数据更新不及时，影响最新月份的计算结果',
            '修复中'
        )

        # 保存测试记录和bug清单
        metrics_calculator.save_test_records()
        metrics_calculator.save_bug_list()

    finally:
        # 关闭数据库连接
        metrics_calculator.close_connection()


if __name__ == "__main__":
    main()
